import torch
import numpy as np
import sys
sys.path.append("..")
from basic_knowledge import d2lzh_pytorch as d2l

#使用Fashion-MNIST数据集。我们将使用多层感知机对图像进行分类
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
#，输入个数为784，输出个数为10。实验中，我们设超参数隐藏单元个数为256。
num_inputs, num_outputs, num_hiddens = 784, 10, 256

W1 = torch.tensor(np.random.normal(0, 0.01, (num_inputs, num_hiddens)), dtype=torch.float)
b1 = torch.zeros(num_hiddens, dtype=torch.float)
W2 = torch.tensor(np.random.normal(0, 0.01, (num_hiddens, num_outputs)), dtype=torch.float)
b2 = torch.zeros(num_outputs, dtype=torch.float)

params = [W1, b1, W2, b2]
for param in params:
    param.requires_grad_(requires_grad=True)

#使用基础的max函数来实现ReLU，而非直接调用relu函数。
def relu(X):
    return torch.max(X, torch.tensor(0.0))
#定义模型，通过view函数将每张原始图像改成长度为num_inputs的向量
# 然后我们实现多层感知机的计算表达式
def net(X):
    X = X.view((-1, num_inputs))
    H = relu(torch.matmul(X, W1) + b1)
    return torch.matmul(H, W2) + b2
#定义损失函数
loss = torch.nn.CrossEntropyLoss()
#训练模型,超参数迭代周期数为5，学习率为100.0。
# 由于原书的mxnet中的SoftmaxCrossEntropyLoss在反向传播的时候相当于沿batch维求和了
# 而PyTorch默认的是求平均，所以用PyTorch计算得到的loss比mxnet小很多
# 大概是maxnet计算得到的1/batch_size这个量级），所以反向传播得到的梯度也小很多，
# 所以为了得到差不多的学习效果，我们把学习率调得成原书的约batch_size倍，
# 原书的学习率为0.5，这里设置成100.0。
# (之所以这么大，应该是因为d2lzh_pytorch里面的sgd函数在更新的时候除以了batch_size
# 其实PyTorch在计算loss的时候已经除过一次了，sgd这里应该不用除了)
num_epochs, lr = 5, 100.0
d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size, params, lr)

